from typing import Optional

from ai_translator_v2.translator import TranslationChain, PDFParser
from ai_translator_v2.translator.writer import Writer


class PDFTranslator:
    def __init__(
            self, model_name: str = 'gpt-3.5-turbo',
    ):
        self.translate_chain = TranslationChain(model_name)
        self.writer = Writer()

    def translate_pdf(self,
                      input_file: str,
                      role: str = '诗人',
                      output_file_path: str = None,
                      output_file_format: str = 'markdown',
                      source_language: str = 'English',
                      target_language: str = 'Chinese',
                      pages: Optional[int] = None):
        book = PDFParser.parse(input_file, pages)
        for page_idx, page in enumerate(book.pages):
            for content_idx, content in enumerate(page.contents):
                translation, status = self.translate_chain.run(content, role, source_language, target_language)
                book.pages[page_idx].contents[content_idx].set_translation(translation, status)
        return self.writer.save_translated_book(book, output_file_format, output_file_path)
